#include <windows.h>
#include <stdio.h>
// #include <string.h>
#include <psapi.h>
#include <stdint.h>
#include <shlwapi.h>

#pragma comment(lib, "User32.lib")
#pragma comment(lib, "Shlwapi.lib")

// Boyer-Moore 搜索算法
void *search_bytes(const uint8_t *haystack, size_t haystack_len,
                   const uint8_t *needle, size_t needle_len)
{
    if (needle_len == 0)
        return (void *)haystack;
    if (haystack_len < needle_len)
        return NULL;
    int n = needle_len - 1;

    // 构建坏字符表
    size_t bad_char[256];
    for (size_t i = 0; i < 256; i++)
    {
        bad_char[i] = needle_len;
    }
    for (size_t i = 0; i < n; i++)
    {
        bad_char[needle[i]] = n - i;
    }

    // 搜索
    size_t offset = 0;
    while (offset <= haystack_len - needle_len)
    {
        int j = n;
        while (j >= 0 && needle[j] == haystack[offset + j])
        {
            j--;
        }
        if (j < 0)
        {
            return (void *)(haystack + offset);
        }
        offset += bad_char[haystack[offset + j]];
    }
    return NULL;
}

// 获取 .text 段的起始地址和大小
void *get_text_section(HMODULE hModule, size_t *text_size)
{
    MODULEINFO moduleInfo;
    GetModuleInformation(GetCurrentProcess(), hModule, &moduleInfo, sizeof(moduleInfo));

    // 遍历节表
    PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)hModule;
    PIMAGE_NT_HEADERS ntHeaders = (PIMAGE_NT_HEADERS)((BYTE *)hModule + dosHeader->e_lfanew);
    PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(ntHeaders);

    for (int i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++, section++)
    {
        if (strcmp((char *)section->Name, ".text") == 0)
        {
            *text_size = section->Misc.VirtualSize;
            return (void *)((BYTE *)hModule + section->VirtualAddress);
        }
    }
    return NULL;
}

// extern int loader_C(void *L) { return 0; };
// extern int loader_Croot(void *L) { return 0; };

BOOL inject_entry_jhqjs(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
{
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
        static wchar_t path[MAX_PATH];
        GetModuleFileNameW(NULL, path, MAX_PATH);
        PWSTR filename = PathFindFileNameW(path);
#ifdef DEBUG
        MessageBoxW(NULL, filename, L"", MB_OK);
#endif
        if (wcscmp(filename, L"dontstarve_steam_x64.exe") != 0 && wcscmp(filename, L"dontstarve_dedicated_server_nullrenderer_x64.exe") != 0)
        {
            return TRUE;
        }
        OutputDebugStringW(L"注入库已加载到目标进程。\n");
#ifdef DEBUG
        MessageBoxW(NULL, L"注入库已加载到目标进程。\n", L"", MB_OK);
#endif

        // 获取 .text 段
        size_t text_size;
        HMODULE base = GetModuleHandle(NULL);
        void *text_start = get_text_section(base, &text_size);
        if (text_start == NULL)
        {
            OutputDebugStringW(L"未找到 .text 段\n");
            MessageBoxW(NULL, L"未找到 .text 段\n", L"", MB_OK);
            return TRUE;
        }
#ifdef DEBUG
        wchar_t buf[1024];
        swprintf_s(buf, 1024, L"地址: %p %zd\n", (uint64_t)text_start - (uint64_t)base, text_size);
        MessageBoxW(NULL, buf, L"", MB_OK);
#endif

        // 计算偏移地址
        const uint8_t needle[] = {0x0f, 0x1f, 0x00, 0x45, 0x33, 0xc0, 0x48, 0x8b, 0xd0, 0x48, 0x8b, 0xcb};
        void *ret = search_bytes((const uint8_t *)text_start, text_size, needle, sizeof(needle));
        if (ret == NULL)
        {
            OutputDebugStringW(L"未找到匹配\n");
            MessageBoxW(NULL, L"未找到匹配\n", L"", MB_OK);
            return TRUE;
        }

        DWORD oldProtect;
        uint64_t addr = (uint64_t)((uint64_t)ret + *((uint32_t *)ret - 1) + 0x10);
#ifdef DEBUG
        swprintf_s(buf, 1024, L"地址: %p %p\n", (uint64_t)ret - (uint64_t)base, (uint64_t)addr - (uint64_t)base);
        MessageBoxW(NULL, buf, L"", MB_OK);
#endif

        VirtualProtect((uint64_t *)addr, sizeof(uint64_t), PAGE_EXECUTE_READWRITE, &oldProtect);

        // 写入新值
        // *(uint64_t *)addr = (uint64_t)loader_C;

        addr += (uint64_t)8;

        *(uint64_t *)addr = (uint64_t)loader_Croot;

        // 恢复保护
        VirtualProtect((uint64_t *)addr, sizeof(uint64_t), oldProtect, &oldProtect);

        OutputDebugStringW(L"内存写入成功。\n");
#ifdef DEBUG
        MessageBoxW(NULL, L"内存写入成功。\n", L"", MB_OK);
#endif
        break;
    }
    return TRUE;
}

int inject_main()
{
    HMODULE mod = LoadLibraryW(L"lapi.dll");
    if (mod == NULL)
    {
        DWORD errorCode = GetLastError();
        LPVOID errorMsg;
        FormatMessage(
            FORMAT_MESSAGE_ALLOCATE_BUFFER |
                FORMAT_MESSAGE_FROM_SYSTEM |
                FORMAT_MESSAGE_IGNORE_INSERTS,
            NULL,
            errorCode,
            MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
            (LPTSTR)&errorMsg,
            0,
            NULL);

        printf("[错误]失败: (代码 %d) %s\n", errorCode, (char *)errorMsg);
        LocalFree(errorMsg);
        MessageBoxW(NULL, L"加载失败\n", L"", MB_OK);
    }
    MessageBoxW(NULL, L"内存写入成功。\n", L"", MB_OK);
    return 0;
}

// // 编译命令
// // cl /utf-8 /DUNICODE /LD /MD /O2 /DLUA_BUILD_AS_DLL /Feinject.dll lapi.c lcode.c ldebug.c ldo.c ldump.c lfunc.c lgc.c llex.c lmem.c lobject.c lopcodes.c lparser.c lstate.c lstring.c ltable.c ltm.c lundump.c lvm.c lzio.c lauxlib.c lbaselib.c ldblib.c liolib.c lmathlib.c loslib.c ltablib.c lstrlib.c loadlib.c linit.c inject.c
// // 编译命令
// // cl /utf-8 /DUNICODE /LD /MD /O2 /DLUA_BUILD_AS_DLL /Fejqsloader.dll lapi.c lcode.c ldebug.c ldo.c ldump.c lfunc.c lgc.c llex.c lmem.c lobject.c lopcodes.c lparser.c lstate.c lstring.c ltable.c ltm.c lundump.c lvm.c lzio.c lauxlib.c lbaselib.c ldblib.c liolib.c lmathlib.c loslib.c ltablib.c lstrlib.c loadlib.c linit.c inject.c
// // cl /utf-8 /DUNICODE /LD /MD /O2  /Fejqsloader.dll inject.c
// // gcc -shared -DUNICODE -o lapi.dll inject.c -lshlwapi -lpsapi